#![feature(try_blocks)]

use futures::FutureExt;

mod rate_limit;
mod slack;
mod symbol_conversion;

#[macro_use]
extern crate sailfish_macros;

use anyhow::Context;
use anyhow::Result;
use clap::Clap;
use futures::{SinkExt, StreamExt};
use itertools::Itertools;
use rate_limit::RateLimit;
use std::collections::HashMap;
use std::io::Read;
use std::path::PathBuf;
use std::sync::Arc;
use strum_macros::Display;
use tokio::stream::StreamExt as TokioStreamExt;


type RelayMessage = (String, chrono::DateTime<chrono::Utc>, String); // (topic, rt, msg)
type LastFeedMap = Arc<chashmap::CHashMap<String, chrono::DateTime<chrono::Utc>>>;

#[derive(Clap, Clone, Debug)]
#[clap(version = "1.0", author = "Donggu Kang")]
struct Opts {
    #[clap(
        long,
        default_value = "../../data/coin2/feed/ops/relay/feed-01.ap-east-1.aws.megaport.json",
        parse(from_os_str)
    )]
    config: PathBuf,
    #[clap(long, default_value = "tcp://*:5124")]
    bind: String,
    #[clap(
        long,
        default_value = "../../data/coin2/feed/subscribers.toml",
        parse(from_os_str)
    )]
    subscriber_toml: PathBuf,
    #[clap(long)]
    debug: bool,
}

#[derive(Clone, Copy, Display)]
enum Decompress {
    None,
    Inflate,
    Gzip,
}

fn unwrap(res: Result<anyhow::Result<()>, tokio::task::JoinError>) -> anyhow::Result<()> {
    res.unwrap_or_else(|e| Err(anyhow::Error::new(e)))
}

// There are many args. We wrap them in this class
// struct SubscriptionRequestArgs<'a> {
//     mea: &'a str,
//     norm_symbol: &'a str,
//     channels: &'a [String],
//     is_every_channel_in_one_subscription: bool
// }

async fn get_subscription_requests(mea: &str, norm_symbol: &str, channels: &[String], is_every_channel_in_one_subscription: bool, sub_template: &str, info: &toml::Value) -> Result<Vec<String>> {
    let symbol_info = symbol_conversion::from_norm(mea, norm_symbol)
    .await
    .with_context(|| {
        format!(
            "could not translate relative-norm '{}' to absolute-norm in mea '{}'",
            mea, norm_symbol
        )
    })?;

    let subscription_requests = {
        let mut reqs = Vec::new();

        // topics = ["futures/trade:BTC-USD-2020..", "futures/book:BTC-USD-2020.."]
        let topics: Vec<String> = channels
            .iter()
            .map(|c| {
                Ok(info["channels"][c]
                    .as_str()
                    .context("subscribers.info 's channels field is not string")?
                    .replace("{}", &symbol_info.feed_native))
            })
            .collect::<Result<Vec<String>>>()?;

        if is_every_channel_in_one_subscription {
            // subscribe symbol.trade, symbol.book in 1 ws

            /*
            {"op": "subscribe", "args": {json_list}}
            {"op": "subscribe", "args": [symbol@trade, symbol@book]}
            */

            let req = sub_template.replace(
                "{json_list}",
                &serde_json::to_string(&topics).with_context(|| {
                    format!(
                        "could not serialize subscription topics to json: {:?}",
                        topics
                    )
                })?,
            );
            reqs.push(req);
        } else {
            for topic in &topics {
                let req = sub_template.replace("{channel}", &topic);
                reqs.push(req);
            }
        }
        reqs
    };
    Ok(subscription_requests)
}

async fn relay_for_symbol(
    mut zmq_tx: tokio::sync::mpsc::Sender<RelayMessage>,
    zmq_topic: String,
    ws_url: String,
    subscription_requests: Vec<String>,
    decompress: Decompress,
    handle_ping_msg: bool,
    rate_limit: Option<RateLimit>,
) -> Result<()> {
    // 0. connect
    // there are two parts.
    // 1. send subscription msg (tx)
    // 2. receive feed msg broadcast (rx)

    // 0. connect

    log::debug!("connect(ws_url={}, topic={})", ws_url, zmq_topic);
    let (conn, _response) = tokio::time::timeout(std::time::Duration::from_secs(10), tokio_tungstenite::connect_async(&ws_url))
        .await?
        .with_context(|| format!("fail to connect to the ws url: {}", ws_url))?;
    let (mut ws_writer, ws_reader) = conn.split();

    // msgs (subscription req / ping-pong) are enqueued to ws_send_tx.
    // there is async thread dedicated to read ws_send_rx and send it to ws.
    let (mut ws_send_tx, mut ws_send_rx) = tokio::sync::mpsc::channel::<String>(100);

    let rate_limiter: Option<_> = if let Some(limit) = rate_limit {
        let limiter = governor::RateLimiter::direct(
            governor::Quota::with_period(limit.duration)
                .context("fail to create rate limiter")?
                .allow_burst(limit.burst),
        );
        Some(limiter)
    } else {
        None
    };

    let mut handles = vec![];

    // 1. a thread responsible for sending msgs to exchange
    handles.push(
        tokio::spawn(async move {
            while let Some(msg_to_send) = ws_send_rx.recv().await {
                ws_writer
                    .send(tungstenite::Message::Text(msg_to_send))
                    .await
                    .context("fail to send a msg to exchange")?;
            }
            Ok::<(), anyhow::Error>(())
        })
        .map(unwrap),
    );

    // enqueue subscription msgs for send (this can happen in parallel to ping/pong)
    {
        let mut ws_send_tx = ws_send_tx.clone();
        handles.push(
            tokio::spawn(async move {
                for req in subscription_requests {
                    if let Some(limiter) = rate_limiter.as_ref() {
                        limiter.until_ready().await;
                    }
                    log::debug!("sending subscription request: {}", req);
                    tokio::time::timeout(std::time::Duration::from_secs(10), ws_send_tx.send(req)).await??;
                }
                Ok::<(), anyhow::Error>(())
            })
            .map(unwrap),
        );
    }

    // 2. receive from exchange
    handles.push(
        tokio::spawn(async move {
            // futures::pin!
            let mut ws_reader = ws_reader.timeout(std::time::Duration::from_secs(5*60));
            while let Some(msg) = ws_reader.try_next().await? {
                let binary_data: Vec<u8>;
                let text_data: String;
                let s: Option<&str> = match msg {
                    Ok(tungstenite::Message::Binary(data)) => {
                        binary_data = match decompress {
                            Decompress::Inflate => miniz_oxide::inflate::decompress_to_vec(&data)
                                .map_err(|e| anyhow::anyhow!("{:?}", e))
                                .context("fail to decode msg with Decompress::Inflate")?,
                            Decompress::Gzip => {
                                let mut buf = Vec::<u8>::new();
                                flate2::bufread::GzDecoder::new(&data[..])
                                    .read_to_end(&mut buf)
                                    .context("fail to decode msg with Decompress::Gzip")?;
                                buf
                            }
                            Decompress::None => data,
                        };
                        Some(std::str::from_utf8(&binary_data).context("not utf-8 string")?)
                    }
                    Ok(tungstenite::Message::Text(s)) => {
                        text_data = s;
                        Some(&text_data)
                    }
                    _ => None,
                };
                if let Some(s) = s {
                    if handle_ping_msg && s.contains("ping") {
                        let _ = ws_send_tx.try_send(s.replace("ping", "pong")); // ignore if full
                        continue;
                    }

                    let rt = chrono::Utc::now();
                    // if full, silently ignore and continue
                    match zmq_tx.try_send((zmq_topic.to_string(), rt, s.to_string())) {
                        Err(e) => {
                            log::error!("{}", e);
                            slack::send(slack_api::chat::PostMessageRequest {
                                text: &format!("[topic={}] probably zmq queue is full", zmq_topic),
                                ..Default::default()
                            })
                            .await?;
                        }
                        _ => {}
                    }
                }
            }

            Ok::<(), anyhow::Error>(())
        })
        .map(unwrap),
    );

    futures::future::try_join_all(handles).await?;
    Ok(())
}

async fn run_relay(opts: Opts, zmq_tx: tokio::sync::mpsc::Sender<RelayMessage>) -> Result<()> {
    let f = std::fs::File::open(&opts.config)?;
    let subscription_config: HashMap<String, Vec<String>> = serde_json::from_reader(f)
        .with_context(|| format!("fail to parse config file: {}", opts.config.display()))?;
    let mut handles = Vec::new();

    let subscriptions_info: toml::Value = toml::from_str(
        &std::fs::read_to_string(&opts.subscriber_toml)
            .with_context(|| format!("fail to read {}", opts.subscriber_toml.display()))?,
    )
    .context("fail to parse subscribers.toml as toml")?;

    for (mear, norm_symbols) in subscription_config {
        let result: Result<()> = try {
            // "incremental", "Futures.Huobi.v1" => "futures_huobi_v1" and "incremental"
            let (recipe, mea) = mear
                .rsplitn(2, ".")
                .collect_tuple()
                .with_context(|| format!("ill-formed mear: {}", mear))?;

            let info = &subscriptions_info[mea.to_lowercase().replace(".", "_")];
            let ws_url = format!(
                "wss://{host}:{port}{path}",
                host = info["host"]
                    .as_str()
                    .with_context(|| format!("host is not str: {}", info["host"]))?,
                port = info["port"]
                    .as_str()
                    .with_context(|| format!("port is not str: {}", info["port"]))?,
                path = info["path"]
                    .as_str()
                    .with_context(|| format!("path is not str: {}", info["path"]))?,
            );

            /* subscription strings
                {"op": "subscribe", "args": {json_list}}

                {"op": "subscribe", "args": ["symbol1@trade", "symbol1@book"]}
                {"op": "subscribe", "args": ["symbol2@trade", "symbol2@book"]}
            */

            // e.g. channels = ["trade", "book_snapshot_5"]
            let channels: Vec<&str> = info["recipes"][recipe]
                .as_array()
                .context("recipes is not array")?
                .iter()
                .map(|v| Ok(v.as_str().context("recipe is not str")?))
                .collect::<Result<Vec<_>>>()?;

            // e.g. sub_template = '{"op": "subscribe", "args": {json_list}}'
            let sub_template = info["subscription_req"]
                .as_str()
                .context("subscription_req is not str")?;
            let is_every_channel_in_one_subscription = sub_template.contains("{json_list}");
            let decompress = match info["decompress"]
                .as_str()
                .context("decompress field is not str")?
            {
                "none" => Decompress::None,
                "gzip" => Decompress::Gzip,
                "deflate" => Decompress::Inflate,
                _ => unreachable!(),
            };
            let handle_ping_msg = info["handle_ping_msg"]
                .as_bool()
                .context("handle_ping_msg field is not bool")?;
            // channels.book_snapshot_5 = "{}@depth5@100ms"

            // read "unauthorized_limit.connection" config
            let mear_connection_limiter = {
                if let Some(limit) = info
                    .get("unauthorized_limit")
                    .and_then(|a| a.get("connection"))
                {
                    let limit = limit
                        .as_str()
                        .context("unauthorized_limit.connection is not str")?;
                    let rate_limit = rate_limit::parse_rate_limit(limit)?;
                    let limiter = governor::RateLimiter::direct(
                        governor::Quota::with_period(rate_limit.duration)
                            .context("fail to create rate limiter")?
                            .allow_burst(rate_limit.burst),
                    );
                    Some(limiter)
                } else {
                    None
                }
            };

            let subscription_rate_limit = {
                if let Some(limit) = info
                    .get("unauthorized_limit")
                    .and_then(|a| a.get("subscription"))
                {
                    let limit = limit
                        .as_str()
                        .context("unauthorized_limit.subscription is not str")?;
                    let rate_limit = rate_limit::parse_rate_limit(limit)?;
                    Some(rate_limit)
                } else {
                    None
                }
            };

            // 1 symbol = 1 ws
            for norm_symbol in &norm_symbols {
                // if disconnects,
                // 1. report to slack
                // 2. sleep
                // 3. try again
                {
                    clone_all::clone_all!(mear, norm_symbol, zmq_tx, ws_url, info);
                    let mea = mea.to_string();
                    let sub_template = sub_template.to_string();
                    let channels:Vec<String> = channels.iter().map(|s| s.to_string()).collect();

                    if let Some(limiter) = mear_connection_limiter.as_ref() {
                        limiter.until_ready().await;
                    }

                    handles.push(
                        tokio::spawn(async move {
                            loop {
                                // // to handle rollover, we resolve subscription requests every connection.
                                let subscription_requests:Vec<String> = get_subscription_requests(&mea, &norm_symbol, &channels[..], is_every_channel_in_one_subscription, &sub_template, &info).await?;
                                let result = relay_for_symbol(
                                    zmq_tx.clone(),
                                    format!("{}:{}", mear, norm_symbol),
                                    ws_url.clone(),
                                    subscription_requests.clone(),
                                    decompress,
                                    handle_ping_msg,
                                    subscription_rate_limit,
                                )
                                .await
                                .with_context(|| "feed ws disconnected");

                                if let Err(e) = result {
                                    slack::send(slack_api::chat::PostMessageRequest {
                                        text: &format!("[{}] [{}:{}] disconnected for error: {:?}\n@donggu", get_hostname(), mear, norm_symbol, e),
                                        channel: "#coin_feed_noti",
                                        link_names: Some(true),
                                        ..Default::default()
                                    })
                                    .await?;
                                }
                            }
                            #[allow(unreachable_code)]
                            Ok::<(), anyhow::Error>(())
                        })
                        .map(unwrap),
                    );
                }

                // wait(ws_connection_interval);
            }
        };
        result.with_context(|| format!("error from mear={} symbols={:?}", mear, norm_symbols))?;
    }

    futures::future::try_join_all(handles).await?;
    Ok(())
}

async fn broadcast(
    opts: Opts,
    mut zmq_rx: tokio::sync::mpsc::Receiver<RelayMessage>,
    last_feed: LastFeedMap,
) -> Result<()> {
    let mut zmq_socket = tmq::publish(&tmq::Context::new())
        .bind(&opts.bind)
        .with_context(|| format!("could not bind to {}", opts.bind))?;
    while let Some((topic, rt, msg)) = zmq_rx.recv().await {
        if opts.debug {
            log::debug!("topic: {}", topic);
            log::debug!("rt: {}", rt);
            log::debug!("msg: {}", msg);
            log::debug!("");
        }
        // let rt = rt.timestamp() * 1_000_000_000 + rt.timestamp_subsec_nanos() as i64;

        last_feed.insert(topic.to_string(), rt);
        let rt = rt.timestamp_nanos().to_string();
        zmq_socket
            .send(vec![&topic, &rt, &msg])
            .await
            .context("fail to zmq_socket.send([topic, rt, msg])")?;
    }

    Ok(())
}

use sailfish::TemplateOnce;
#[derive(TemplateOnce)]
#[template(path = "dashboard.stpl")]
struct DashboardTemplate<'a> {
    host: &'a str,
    bind: &'a str,
    config: &'a std::path::Path,
    start: String,
    uptime: String,
    subscription_config: &'a HashMap<String, Vec<String>>,
    last_feed: LastFeedMap,
}

#[rocket::get("/")]
async fn get_dashboard(
    start: rocket::State<'_, chrono::DateTime<chrono::Utc>>,
    opts: rocket::State<'_, Opts>,
    subscription_config: rocket::State<'_, HashMap<String, Vec<String>>>,
    last_feed: rocket::State<'_, LastFeedMap>,
) -> rocket::response::content::Html<String> {
    let host = hostname::get().unwrap_or(std::ffi::OsString::new());
    let template = DashboardTemplate {
        host: host.to_str().unwrap_or(""),
        bind: &opts.bind,
        config: &opts.config,
        start: start.format("%Y-%m-%d %H:%M").to_string(),
        uptime: humantime::format_duration((chrono::Utc::now() - *start).to_std().unwrap())
            .to_string(),
        subscription_config: &subscription_config,
        last_feed: last_feed.clone(),
    };

    let html = template.render_once().unwrap();
    rocket::response::content::Html(html)
}

fn get_hostname() -> String {
    hostname::get().unwrap_or(std::ffi::OsString::new()).to_string_lossy().into_owned()
}

async fn main_impl() -> Result<()> {
    let opts: Opts = Opts::parse();
    log::info!("{:?}", opts);

    let (zmq_tx, zmq_rx) = tokio::sync::mpsc::channel(1000);
    let last_feed = LastFeedMap::default();

    slack::send(slack_api::chat::PostMessageRequest {
        text: &format!("[{}] feed-relay started.", get_hostname()),
        // channel: "#coin_feed_noti",
        channel: "#coin_feed_noti",
        ..Default::default()
    })
    .await?;

    let mut handles = vec![];
    // 1. subscribe & receive feed
    {
        let opts = opts.clone();
        handles.push(
            tokio::spawn(async move {
                run_relay(opts, zmq_tx).await?;
                Ok::<(), anyhow::Error>(())
            })
            .map(unwrap),
        );
    } // JoinHandle::Output = Result<T, JoinError> = Result<Result<T', Err>, JoinError>

    // 2. broadcast feed
    {
        clone_all::clone_all!(opts, last_feed);
        handles.push(
            tokio::spawn(async move {
                broadcast(opts, zmq_rx, last_feed).await?;
                Ok::<(), anyhow::Error>(())
            })
            .map(unwrap),
        );
    }

    // 3. web dashboard
    handles.push(
        tokio::spawn(async move {
            let f = std::fs::File::open(&opts.config)?;
            let subscription_config: HashMap<String, Vec<String>> = serde_json::from_reader(f)
                .with_context(|| {
                    format!("fail to parse config file: {}", &opts.config.display())
                })?;

            rocket::ignite()
                .manage(chrono::Utc::now())
                .manage(opts)
                .manage(subscription_config)
                .manage(last_feed)
                .mount("/", rocket::routes![get_dashboard])
                .launch()
                .await
                .unwrap();
            Ok::<(), anyhow::Error>(())
        })
        .map(unwrap),
    );

    futures::future::try_join_all(handles).await?;
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();
    let result = main_impl().await;
    if let Err(e) = result {
        log::error!("{:?}", e);
        slack::send(slack_api::chat::PostMessageRequest {
            text: &format!("[{}] [error] feed-relay terminates: {:?}", get_hostname(), e),
            // channel: "#coin_feed_noti",
            channel: "#coin_feed_noti",
            ..Default::default()
        })
        .await?;
    }
    Ok(())
}
